"""
量化交易指标计算工具模块
"""
import numpy as np
import pandas as pd
from typing import Union, Optional, Dict, Any, List, Tuple
from decimal import Decimal


class Indicators:
    """
    通用技术指标计算工具类
    提供各种常用和自定义指标的计算方法
    """
    
    @staticmethod
    def custom_sar(
        high: pd.Series,
        low: pd.Series,
        close: pd.Series,
        start: float = 0.02,  # N2 / 100
        step: float = 0.02,   # N3 / 100
        limit: float = 0.2    # N4 / 100
    ) -> pd.Series:
        """
        自定义SAR（抛物线转向）指标，根据AICOIN研究院的实现

        Args:
            high: 最高价序列
            low: 最低价序列
            close: 收盘价序列
            start: 初始加速因子 (默认 0.02)
            step: 加速因子步长 (默认 0.02)
            limit: 加速因子最大值 (默认 0.2)

        Returns:
            pd.Series: SAR值的序列
        """
        # 确保输入数据是浮点型
        high = high.astype(float)
        low = low.astype(float)
        close = close.astype(float)
        
        # 初始化结果数组
        sar = np.zeros(len(close))
        sar.fill(np.nan)
        
        # 初始化辅助变量
        af = np.zeros(len(close))  # 加速因子
        max_min = np.zeros(len(close))  # 极值点
        up = np.zeros(len(close), dtype=bool)  # 趋势方向
        is_first_trend_bar = np.zeros(len(close), dtype=bool)  # 是否是趋势的第一个柱
        
        # 至少需要2个数据点才能开始
        if len(close) < 2:
            return pd.Series(sar)
            
        # 初始化第一个SAR值
        if close.iloc[1] > close.iloc[0]:
            up[1] = True
            max_min[1] = high.iloc[1]
            sar[1] = low.iloc[0]
        else:
            up[1] = False
            max_min[1] = low.iloc[1]
            sar[1] = high.iloc[0]
        
        af[1] = start
        is_first_trend_bar[1] = True
        
        # 逐点计算SAR
        for i in range(2, len(close)):
            # 延续上一个周期的状态
            up[i] = up[i-1]
            max_min[i] = max_min[i-1]
            af[i] = af[i-1]
            is_first_trend_bar[i] = False
            
            # 临时保存上一个SAR值
            tmp_sar = sar[i-1]
            
            # 计算当前的SAR值
            sar[i] = tmp_sar + af[i] * (max_min[i] - tmp_sar)
            
            # 检查是否需要反转
            if up[i]:
                if sar[i] > low.iloc[i]:
                    # 反转为向下趋势
                    is_first_trend_bar[i] = True
                    up[i] = False
                    sar[i] = max(high.iloc[i], max_min[i])
                    max_min[i] = low.iloc[i]
                    af[i] = start
                else:
                    # 继续向上趋势
                    if high.iloc[i] > max_min[i]:
                        max_min[i] = high.iloc[i]
                        af[i] = min(af[i] + step, limit)
            else:
                if sar[i] < high.iloc[i]:
                    # 反转为向上趋势
                    is_first_trend_bar[i] = True
                    up[i] = True
                    sar[i] = min(low.iloc[i], max_min[i])
                    max_min[i] = high.iloc[i]
                    af[i] = start
                else:
                    # 继续向下趋势
                    if low.iloc[i] < max_min[i]:
                        max_min[i] = low.iloc[i]
                        af[i] = min(af[i] + step, limit)
            
            # 对SAR值进行限制
            if up[i]:
                if i >= 3:
                    sar[i] = min(sar[i], low.iloc[i-1], low.iloc[i-2])
            else:
                if i >= 3:
                    sar[i] = max(sar[i], high.iloc[i-1], high.iloc[i-2])
        
        return pd.Series(sar, index=close.index)

    @staticmethod
    def aicoin_sar(df: pd.DataFrame, n2: float = 2.0, n3: float = 2.0, n4: float = 20.0) -> pd.Series:
        """
        AICOIN研究院实现的SAR指标，直接使用原始参数

        Args:
            df: 包含high, low, close列的DataFrame
            n2: 初始加速因子百分比 (默认 2.0)
            n3: 加速因子步长百分比 (默认 2.0)
            n4: 加速因子最大值百分比 (默认 20.0)

        Returns:
            pd.Series: SAR值的序列
        """
        # 确保数据是浮点型
        high = df['high'].astype(float)
        low = df['low'].astype(float)
        close = df['close'].astype(float)
        
        start = float(n2) / 100.0
        step = float(n3) / 100.0
        limit = float(n4) / 100.0
        
        return Indicators.custom_sar(high, low, close, start, step, limit)
        
    @staticmethod
    def supertrend(
        high: pd.Series,
        low: pd.Series,
        close: pd.Series,
        periods: int = 10,
        multiplier: float = 3.9,
        source: str = 'hl2'
    ) -> Dict[str, pd.Series]:
        """
        计算SuperTrend指标
        
        Args:
            high: 最高价序列
            low: 最低价序列
            close: 收盘价序列
            periods: ATR周期 (默认 10)
            multiplier: ATR乘数 (默认 3.9)
            source: 价格源类型 ('hl2'=高低价平均, 'ohlc4'=开高低收平均)
            
        Returns:
            Dict[str, pd.Series]: 包含以下键的字典:
                - 'trend': 趋势方向 (1=上涨, -1=下跌)
                - 'up': 上轨线
                - 'down': 下轨线
                - 'buy_signal': 买入信号
                - 'sell_signal': 卖出信号
        """
        # 确保输入数据是浮点型
        high = high.astype(float)
        low = low.astype(float)
        close = close.astype(float)
        multiplier = float(multiplier)
        
        # 计算价格源
        if source == 'hl2':
            price = (high + low) / 2
        elif source == 'ohlc4':
            price = (high + low + close + close) / 4  # 收盘价权重加倍
        else:
            price = close  # 默认使用收盘价
            
        # 计算ATR
        tr = pd.Series(np.zeros(len(close)))
        for i in range(len(close)):
            if i == 0:
                tr.iloc[i] = float(high.iloc[i] - low.iloc[i])
            else:
                tr.iloc[i] = float(max(
                    high.iloc[i] - low.iloc[i],
                    abs(high.iloc[i] - close.iloc[i-1]),
                    abs(low.iloc[i] - close.iloc[i-1])
                ))
        
        # 使用SMMA计算ATR (等同于EMA但初始化不同)
        atr = pd.Series(np.zeros(len(close)))
        for i in range(len(close)):
            if i == 0:
                atr.iloc[i] = float(tr.iloc[i])
            else:
                if i < periods:
                    atr.iloc[i] = (atr.iloc[i-1] * (i) + tr.iloc[i]) / (i + 1)
                else:
                    atr.iloc[i] = (atr.iloc[i-1] * (periods - 1) + tr.iloc[i]) / periods
        
        # 计算上下轨
        up = pd.Series(np.zeros(len(close)))
        dn = pd.Series(np.zeros(len(close)))
        
        for i in range(len(close)):
            up.iloc[i] = float(price.iloc[i] - (multiplier * atr.iloc[i]))
            dn.iloc[i] = float(price.iloc[i] + (multiplier * atr.iloc[i]))
        
        # 初始化趋势
        trend = pd.Series(np.ones(len(close)))
        
        # 计算SuperTrend
        for i in range(1, len(close)):
            # 上轨线计算
            if close.iloc[i-1] > up.iloc[i-1]:
                up.iloc[i] = max(up.iloc[i], up.iloc[i-1])
            
            # 下轨线计算
            if close.iloc[i-1] < dn.iloc[i-1]:
                dn.iloc[i] = min(dn.iloc[i], dn.iloc[i-1])
            
            # 趋势判断
            if close.iloc[i] > dn.iloc[i-1]:
                trend.iloc[i] = 1
            elif close.iloc[i] < up.iloc[i-1]:
                trend.iloc[i] = -1
            else:
                trend.iloc[i] = trend.iloc[i-1]
                
                # 根据趋势调整上下轨
                if trend.iloc[i] == 1 and up.iloc[i] < up.iloc[i-1]:
                    up.iloc[i] = up.iloc[i-1]
                if trend.iloc[i] == -1 and dn.iloc[i] > dn.iloc[i-1]:
                    dn.iloc[i] = dn.iloc[i-1]
        
        # 生成买卖信号
        buy_signal = (trend == 1) & (trend.shift(1) == -1)
        sell_signal = (trend == -1) & (trend.shift(1) == 1)
        
        return {
            'trend': pd.Series(trend, index=close.index),
            'up': pd.Series(up, index=close.index),
            'down': pd.Series(dn, index=close.index),
            'buy_signal': pd.Series(buy_signal, index=close.index),
            'sell_signal': pd.Series(sell_signal, index=close.index)
        } 